from dflow import config, s3_config
config["host"] = "http://39.106.93.187:32746"
config["k8s_api_server"] = "https://101.200.186.161:6443"
config["token"] = "eyJhbGciOiJSUzI1NiIsImtpZCI6IlhMRGZjbnNRemE4RGQyUXRMZG1MX3NXeG5TMzlQTnhnSHZkS1lGM25SODAifQ.eyJpc3MiOiJrdWJlcm5ldGVzL3NlcnZpY2VhY2NvdW50Iiwia3ViZXJuZXRlcy5pby9zZXJ2aWNlYWNjb3VudC9uYW1lc3BhY2UiOiJhcmdvIiwia3ViZXJuZXRlcy5pby9zZXJ2aWNlYWNjb3VudC9zZWNyZXQubmFtZSI6ImFyZ28tdG9rZW4tOGY4djkiLCJrdWJlcm5ldGVzLmlvL3NlcnZpY2VhY2NvdW50L3NlcnZpY2UtYWNjb3VudC5uYW1lIjoiYXJnbyIsImt1YmVybmV0ZXMuaW8vc2VydmljZWFjY291bnQvc2VydmljZS1hY2NvdW50LnVpZCI6IjBhNzI1N2JhLWZkZWQtNGI2OS05YWU2LTZhY2U0M2UxNjdlNiIsInN1YiI6InN5c3RlbTpzZXJ2aWNlYWNjb3VudDphcmdvOmFyZ28ifQ.ocRNsp_sdjM7dKpg_rYwPOKATslbDebDo807rkaVcqOScFKl11tqHbCwsekq4BZSEw-MaZjWAVsdE5jgvqIggp2qczx1QPMkBGQzkPwR4h7HbxYPUKIHkjlcvtsl06jf7urfzARUiD_UTahEFQLkUeN800Qblp-zMFBNLF2Y7_wW867drmQxynG1ssQ8agdu7yDZpwwJz-qzMMZsuZ7QNtL0pPQP2Iw_5C6jlos-Al1m3bgUJ5phh7yt-PBqwnMwEPaZWkVy9zFJc5t4J0jFc9nRnrR8fEd_OgJTTgQjHxS2DyXj9ZRUlGJA-tmHGqIJ7nuScv3lKwbb8TkeABq5DA"
s3_config["endpoint"] = "39.106.93.187:30900"

from dflow import (
    InputParameter,
    OutputParameter,
    InputArtifact,
    OutputArtifact,
    upload_artifact,
    Workflow,
    Step
)
from dflow.plugins.lebesgue import LebesgueContext

from dflow.python import (
    PythonOPTemplate,
    OP,
    OPIO,
    OPIOSign,
    Artifact,
)

import os, time
from typing import List
from pathlib import Path

from ase.optimize import BFGS
from ase.calculators.emt import EMT
from ase.io import read, write
from ase.neb import NEB

def structure_optimization(structure_file: Path) -> Path:
    # Load structure
    structure = read(structure_file)

    # Set up the calculator
    structure.set_calculator(EMT())

    # Optimize structure
    opt = BFGS(structure)
    opt.run(fmax=0.05)

    # Write optimized structure to file
    optimized_structure_file = structure_file.with_stem(structure_file.stem + '_optimized')
    write(optimized_structure_file, structure)

    return optimized_structure_file

def neb_calculation(initial_structure: Path, final_structure: Path, images: int) -> Path:
    # Load initial and final structures
    initial = read(initial_structure)
    final = read(final_structure)

    # Make a band consisting of 'images' number of images
    images = [initial]
    images += [initial.copy() for i in range(images)]
    images += [final]

    # Set up NEB and the calculator
    neb = NEB(images)
    neb.interpolate()
    for image in images:
        image.set_calculator(EMT())

    # Optimize NEB path
    opt = BFGS(neb, trajectory='neb.traj')
    opt.run(fmax=0.05)

    return Path('neb.traj')

class StructureOptimizationOP(OP):
    @classmethod
    def get_input_sign(cls) -> OPIOSign:
        return OPIOSign({
            "initial_structure": Artifact(Path),
            "final_structure": Artifact(Path)
        })

    @classmethod
    def get_output_sign(cls) -> OPIOSign:
        return OPIOSign({
            "optimized_initial_structure": Artifact(Path),
            "optimized_final_structure": Artifact(Path),
        })

    @OP.exec_sign_check
    def execute(self, op_in: OPIO) -> OPIO:
        initial_structure = op_in["initial_structure"]
        final_structure = op_in["final_structure"]

        optimized_initial_structure = structure_optimization(initial_structure)
        optimized_final_structure = structure_optimization(final_structure)

        op_out = {
            "optimized_initial_structure": optimized_initial_structure,
            "optimized_final_structure": optimized_final_structure,
        }
        return op_out

class NEBCalculationOP(OP):
    @classmethod
    def get_input_sign(cls) -> OPIOSign:
        return OPIOSign({
            "neb_initial_structure": Artifact(Path),
            "neb_final_structure": Artifact(Path),
            "images": Parameter(int)
        })

    @classmethod
    def get_output_sign(cls) -> OPIOSign:
        return OPIOSign({
            "neb_path": Artifact(Path)
        })

    @OP.exec_sign_check
    def execute(self, op_in: OPIO) -> OPIO:
        neb_initial_structure = op_in["neb_initial_structure"]
        neb_final_structure = op_in["neb_final_structure"]
        images = op_in["images"]

        neb_path = neb_calculation(neb_initial_structure, neb_final_structure, images)

        op_out = {
            "neb_path": neb_path
        }
        return op_out

def main():
    lebesgue_context = LebesgueContext(
        username="your username",
        password="your password",
        executor="lebesgue_v2",
        extra={
                "scass_type":"c16_m32_cpu",
                "program_id":your program ID,
                "job_type":"container",
                "template_cover_cmd_escape_bug":"True"
        },
    )

artifact0 = upload_artifact("./initial_stru")
artifact1 = upload_artifact("./final_stru")

# Create steps
structure_optimization_step = Step(
    name="structure_optimization",
    template=PythonOPTemplate(StructureOptimization, image="dp-harbor-registry.cn-zhangjiakou.cr.aliyuncs.com/deepmd/vasp:5.4.4"),
    artifacts={"initial_structure":artifact0,"final_structure":artifact1}
)

neb_calculation_step = Step(
    name="neb_calculation",
    template=PythonOPTemplate(NEBCalculation, image="dp-harbor-registry.cn-zhangjiakou.cr.aliyuncs.com/deepmd/vasp:5.4.4"),
    parameters={
        "neb_initial_structure": structure_optimization_step.outputs["optimized_initial_structure"],
        "neb_final_structure": structure_optimization_step.outputs["optimized_final_structure"],
        "images": 3
    }
)

# Create workflow
wf = Workflow(name="neb_workflow")
wf.add(structure_optimization_step)
wf.add(neb_calculation_step)

# Submit workflow
wf.submit()
